package LeetCode.month04.day105;

/**
 * @author 16778
 * @date 2025/11/15 20:47
 * @description
 */
public class No3234 {
    public int numberOfSubstrings(String s) {
        int n = s.length();
        int[] pos0 = new int[n + 1]; // 0 的下标
        pos0[0] = -1; // 加个 -1 哨兵，方便处理 cnt0 达到最大时的计数
        int size = 1;
        int total1 = 0; // [0,r] 中的 1 的个数
        int ans = 0;

        for (int r = 0; r < n; r++) {
            if (s.charAt(r) == '0') {
                pos0[size++] = r; // 记录 0 的下标
            } else {
                total1++;
                ans += r - pos0[size - 1]; // 单独计算不含 0 的子串个数
            }

            // 倒着遍历 pos0，那么 cnt0 = size - i
            for (int i = size - 1; i > 0 && (size - i) * (size - i) <= total1; i--) {
                int p = pos0[i - 1];
                int q = pos0[i];
                int cnt0 = size - i;
                int cnt1 = r - q + 1 - cnt0; // [q,r] 中的 1 的个数 = [q,r] 的长度 - cnt0
                ans += Math.max(q - Math.max(cnt0 * cnt0 - cnt1, 0) - p, 0);
            }
        }

        return ans;
    }
}
